#include "evdeploy/cv/detection/yolov7_prepostprocessor.h"

namespace ev
{
    namespace vision
    {
        YOLOv7Preprocessor::YOLOv7Preprocessor()
        {
        }

        YOLOv7Preprocessor::~YOLOv7Preprocessor()
        {
        }

        float YOLOv7Preprocessor::GetScale()
        {
            return m_scale;
        }

        EVStatus YOLOv7Preprocessor::Run(cv::Mat &in_mat, cv::Mat &out_mat, int out_width, bool normalize, bool bgr2rgb)
        {
            REC_TIME(t0);
            if (out_width < 10)
            {
                EVLOG(ERROR) << "input size too small:" << out_width;
                return EV_FAIL;
            }
            m_scale = std::min(out_width / static_cast<float>(in_mat.rows), out_width / static_cast<float>(in_mat.cols));
            cv::Size new_size = cv::Size{in_mat.cols * m_scale, in_mat.rows * m_scale};
            cv::Mat resized_mat;
            cv::resize(in_mat, resized_mat, new_size);

            cv::Mat board = cv::Mat(cv::Size(out_width, out_width), CV_8UC3, cv::Scalar(114, 114, 114));

            resized_mat.copyTo(board(cv::Rect{0, 0, resized_mat.cols, resized_mat.rows}));

            if (bgr2rgb)
            {
                cv::cvtColor(board, board, cv::COLOR_BGR2RGB);
            }

            if (normalize)
            {
                board.convertTo(board, CV_32F, 1 / 255.);
            }
            else
            {
                board.convertTo(board, CV_32F);
            }

            out_mat = board.clone();
            std::vector<cv::Mat> chw_wrappers;
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data + sizeof(float) * out_width * out_width);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data + 2 * sizeof(float) * out_width * out_width);
            cv::split(board, chw_wrappers);

            REC_TIME(t1);
            EVLOG(INFO) << "YOLOv7Preprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }

        YOLOv7Postprocessor::YOLOv7Postprocessor()
        {
        }

        YOLOv7Postprocessor::~YOLOv7Postprocessor()
        {
        }

        EVStatus YOLOv7Postprocessor::Run(EVMatData *out, std::vector<BoxInfo> &objects, float scale, float thresh, const int img_w, const int img_h, float nms_thresh, bool class_aware)
        {

            REC_TIME(t0);
            std::vector<BoxInfo> proposals;

            // need to know the net output shape
            int box_num = out->dims[1];
            // need to know the output class num
            int class_num = out->dims[2] - 5;
            float *buffer = (float *)out->data;
            for (int i = 0; i < box_num; ++i)
            {
                int index = i * (class_num + 5);
                // EVLOG(INFO) << buffer[index+4];
                if (buffer[index + 4] > thresh) // det_thresh 0.2
                {
                    float x = buffer[index] / scale;
                    float y = buffer[index + 1] / scale;
                    float w = buffer[index + 2] / scale;
                    float h = buffer[index + 3] / scale;
                    float *max_cls_pos = std::max_element(buffer + index + 5, buffer + index + 5 + class_num);

                    if ((*max_cls_pos) * buffer[index + 4] > 0.2)
                    {
                        cv::Rect box{x - w / 2, y - h / 2, w, h};
                        box = box & cv::Rect(0, 0, img_w - 1, img_h - 1);
                        if (box.area() > 0)
                        {
                            BoxInfo box_info = {box.x, box.y, box.x + box.width, box.y + box.height, (*max_cls_pos) * buffer[index + 4], max_cls_pos - (buffer + index + 5)};
                            objects.push_back(box_info);
                        }
                    }
                }
            }

            NMS(objects, nms_thresh, class_aware);
            REC_TIME(t1);
            EVLOG(INFO) << "YOLOv7Postprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }

    } // namespace vision
} // namespace ev
